package com.example.controller;

import com.gearwenxin.client.ernie.ErnieBot4Client;
import com.gearwenxin.client.stable.StableDiffusionXLClient;
import com.gearwenxin.entity.enums.SamplerType;
import com.gearwenxin.entity.request.ImageBaseRequest;
import com.gearwenxin.entity.response.ChatResponse;
import com.gearwenxin.entity.response.ImageResponse;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Mono;

import javax.annotation.Resource;

@RestController
@RequestMapping("/ai")
public class ChatController {

    // 要调用的模型的客户端（示例为文心4.0）
    @Resource
    private ErnieBot4Client ernieBot4Client;
    @Resource
    private StableDiffusionXLClient stableDiffusionXLClient;
    @Autowired
    private WebClient.Builder webClientBuilder;

    // 单次对话
    @PostMapping("/chat")
    public Mono<ChatResponse> chatSingle(String msg) {
        return ernieBot4Client.chatSingle(msg);
    }

    // 流式返回，连续对话（SSE形式）
    @GetMapping(value = "/stream/chats", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter chatSingleSSE(@RequestParam String msg) {
        String chatUID = "test-user-1001";
        SseEmitter emitter = new SseEmitter();
        ernieBot4Client.chatContOfStream(msg, chatUID).subscribe(response -> {
            try {
                emitter.send(SseEmitter.event().data(response.getResult()));
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
        }, emitter::completeWithError, emitter::complete);
        return emitter;
    }

    // Prompt模板
/*  @PostMapping("/prompt")
  public Mono<PromptResponse> chatSingle() {
      Map<String, String> map = new HashMap<>();
      map.put("article", "我看见过波澜壮阔的大海，玩赏过水平如镜的西湖，却从没看见过漓江这样的水。漓江的水真静啊，静得让你感觉不到它在流动。");
      map.put("number", "20");
      PromptRequest promptRequest = new PromptRequest();
      promptRequest.setId("1234");
      promptRequest.setParamMap(map);

      return promptBotClient.chatPrompt(promptRequest);
  }*/

    // 文生图
    @PostMapping("/image")
    public Mono<ImageResponse> chatImage() {
        ImageBaseRequest imageBaseRequest = ImageBaseRequest.builder()
                // 提示词
                .prompt("一个头发中分背带裤的人")
                // 大小
                .size("1024x1024")
                // 反向提示词（不包含什么）
                .negativePrompt("鸡")
                // 生成图片数量（1-4）
                .n(1)
                // 迭代轮次（10-50）
                .steps(10)
                // 采样方式
                .samplerIndex(SamplerType.Euler_A.getValue())
                .userId("1001")
                .build();


        // 使用配置好的WebClient
        WebClient webClient = webClientBuilder.build();
        Mono<ImageResponse> imageResponseMono = stableDiffusionXLClient.chatImage(imageBaseRequest);

        imageResponseMono.doOnNext(imageResponse -> {
            System.out.println("imageResponse = " + imageResponse);
        }).subscribe();

        return imageResponseMono;
    }

}